package com.anan;

import javax.imageio.ImageIO;
import javax.swing.*;
import java.awt.*;
import java.awt.event.MouseAdapter;
import java.awt.event.MouseEvent;
import java.awt.event.MouseMotionAdapter;
import java.awt.image.BufferedImage;
import java.io.*;
import java.nio.charset.Charset;
import java.nio.file.Files;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;

/**
 * project_name: easyai-labelimg
 * Author: Anan
 * time: 2025/2/22 12:10
 */
public class YoloLabeler extends JFrame {
    private File currentFolder;  // 选择的文件目录
    private List<File> imageFiles = new ArrayList<>();
    private int currentIndex = 0;   // 选择的图片索引

    private JLabel statusBar = new JLabel("    Ready");  //
    private ImagePanel imagePanel = new ImagePanel();  // 图片显示区域


    private JComboBox<String> classCombo = new JComboBox<>();  // 类别选择
    private Point startPoint;  // 目标框开始坐标x，y
    private Rectangle currentRect;  // 当前绘制的目标框矩形


    public YoloLabeler() {
        initUI();
        setVisible(true);
    }

    private void initUI() {
        try {
            String classnamesPath = System.getProperty("user.dir")+File.separator+"classnames.txt";
            File file = new File(classnamesPath);
            List<String> result = Files.readAllLines(file.toPath(), Charset.defaultCharset());
            result.forEach(line -> {
                classCombo.addItem(line.split(" ")[0]);
            });
        }catch (Exception e) {
            System.err.println("请检查设置 classnamesPath classnames.txt 文件的路径");
            e.printStackTrace();
        }




        setTitle("YOLO Labeler");
        setSize(1200, 800);
        setDefaultCloseOperation(EXIT_ON_CLOSE);

        // 工具栏
        JToolBar toolBar = new JToolBar();
        JButton openBtn = new JButton("打开文件夹");
        openBtn.setPreferredSize(new Dimension(100, 20)); // 设置宽度
        openBtn.addActionListener(e -> openFolder());
        toolBar.add(Box.createHorizontalStrut(10)); // 添加间距
        toolBar.add(openBtn);

        // 类别选择
        toolBar.addSeparator();
        JLabel jLabel = new JLabel("Class:");
        jLabel.setFont(new Font("Arial", Font.BOLD, 15)); // 设置字体
//        jLabel.setPreferredSize(new Dimension(100, 20)); // 设置宽度
        toolBar.add(jLabel);
        toolBar.add(classCombo);
        toolBar.add(Box.createHorizontalStrut(10)); // 添加间距

        // 导航按钮
        JButton prevBtn = new JButton("上一张");
        prevBtn.setPreferredSize(new Dimension(100, 20)); // 设置宽度
        prevBtn.addActionListener(e -> showImage(currentIndex - 1));
        toolBar.add(Box.createHorizontalStrut(10)); // 添加间距
        toolBar.add(prevBtn);

        JButton nextBtn = new JButton("下一张");
        nextBtn.setPreferredSize(new Dimension(100, 20));
        nextBtn.addActionListener(e -> showImage(currentIndex + 1));
        toolBar.add(Box.createHorizontalStrut(10)); // 添加间距
        toolBar.add(nextBtn);

        JButton saveBoxBtn = new JButton("保存");
        saveBoxBtn.setPreferredSize(new Dimension(100, 20));
        saveBoxBtn.addActionListener(e -> saveAnnotations());
        toolBar.add(Box.createHorizontalStrut(10)); // 添加间距
        toolBar.add(saveBoxBtn);

        JButton deleteBoxBtn = new JButton("删除矩形框");
        deleteBoxBtn.setPreferredSize(new Dimension(100, 20));
        deleteBoxBtn.addActionListener(e -> deleteBox());
        toolBar.add(Box.createHorizontalStrut(10)); // 添加间距
        toolBar.add(deleteBoxBtn);

        statusBar.setPreferredSize(new Dimension(0, 40));
        statusBar.setFont(new Font("Arial", Font.BOLD, 20));
        // 主布局
        add(toolBar, BorderLayout.NORTH);
        add(new JScrollPane(imagePanel), BorderLayout.CENTER);
        add(statusBar, BorderLayout.SOUTH);


        // 鼠标事件
        imagePanel.addMouseListener(new MouseAdapter() {
            public void mousePressed(MouseEvent e) {
                startPoint = e.getPoint();
                imagePanel.tempRect.x = e.getX();
                imagePanel.tempRect.y = e.getY();

                currentRect = new Rectangle(startPoint.x, startPoint.y, 0, 0);
            }

            public void mouseReleased(MouseEvent e) {
                if (currentRect.width > 10 && currentRect.height > 10) { // 最小尺寸限制
                    Point endPoint = e.getPoint();
                    currentRect.setSize(endPoint.x - startPoint.x, endPoint.y - startPoint.y);

                    // 转换为图片实际坐标
                    Rectangle actualRect = imagePanel.getActualRect(currentRect);
                    if (actualRect != null) {
                        String className = (String) classCombo.getSelectedItem();
                        imagePanel.annotations.add(new YoloAnnotation(
                                className,
                                actualRect.x,
                                actualRect.y,
                                actualRect.width,
                                actualRect.height
                        ));
                        imagePanel.repaint();
                    }
                    imagePanel.setTempRect(new Rectangle());
                }
            }
        });

        imagePanel.addMouseMotionListener(new MouseMotionAdapter() {
            public void mouseDragged(MouseEvent e) {
                if (startPoint != null) {
                    imagePanel.tempRect.width = e.getX() - startPoint.x;
                    imagePanel.tempRect.height = e.getY() - startPoint.y;
                    currentRect.setSize(
                            e.getX() - startPoint.x,
                            e.getY() - startPoint.y
                    );
                    imagePanel.setTempRect(currentRect);
                    imagePanel.repaint();
                }
            }
        });
    }

    private void openFolder() {
        JFileChooser chooser = new JFileChooser();  // 创建
        chooser.setFileSelectionMode(JFileChooser.DIRECTORIES_ONLY);
        if (chooser.showOpenDialog(this) == JFileChooser.APPROVE_OPTION) {
            currentFolder = chooser.getSelectedFile();
            loadImageFiles();
            showImage(0);
        }
    }

    private void loadImageFiles() {
        imageFiles.clear();
        File[] files = currentFolder.listFiles((dir, name) ->
                name.matches(".*\\.(jpg|png|jpeg|bmp)"));
        if (files != null) {
            imageFiles.addAll(Arrays.asList(files));
        }
    }

    private void showImage(int index) {
        if (index >= 0 && index < imageFiles.size()) {
            try {
                currentIndex = index;
                BufferedImage image = ImageIO.read(imageFiles.get(index));
                imagePanel.setImage(image);
                loadAnnotations();
                updateStatus();
            } catch (IOException ex) {
                ex.printStackTrace();
            }
        }
    }

    private void loadAnnotations() {
        imagePanel.annotations.clear();
        File txtFile = getAnnotationFile();
        if (txtFile.exists()) {
            try (BufferedReader br = new BufferedReader(new FileReader(txtFile))) {
                String line;
                while ((line = br.readLine()) != null) {
                    String[] parts = line.split(" ");
                    if (parts.length >= 5) {
                        imagePanel.annotations.add(new YoloAnnotation(
                                parts[0],
                                Integer.parseInt(parts[1]),
                                Integer.parseInt(parts[2]),
                                Integer.parseInt(parts[3]),
                                Integer.parseInt(parts[4])
                        ));
                    }
                }
            } catch (Exception ex) {
                ex.printStackTrace();
            }
        }
    }

    private void saveAnnotations() {
        File txtFile = getAnnotationFile();
        try (PrintWriter pw = new PrintWriter(txtFile)) {
            for (YoloAnnotation ann : imagePanel.annotations) {
                pw.println(String.format("%s %d %d %d %d", ann.classId, ann.x, ann.y, ann.width, ann.height));
            }
        } catch (Exception ex) {
            ex.printStackTrace();
        }
    }

    private File getAnnotationFile() {
        String imgPath = imageFiles.get(currentIndex).getAbsolutePath();
        return new File(imgPath.substring(0, imgPath.lastIndexOf('.')) + ".txt");
    }

    private void deleteBox() {


        if(imagePanel.annotations.size()>0){
            imagePanel.annotations.remove(imagePanel.annotations.size()-1);
            imagePanel.repaint();
        }


    }

    private void updateStatus() {
        statusBar.setText("    Image: " + (currentIndex + 1) + "/" + imageFiles.size()+"  "+imageFiles.get(currentIndex).getAbsolutePath());
    }

}
